

%% data, parameter estimates, and simulated outcomes

load('voting_stage.mat')
load('campaign_stage.mat')
load('coalition_stage_p1.mat')
load('coalition_stage.mat')


%% district-level counterfactual election results (table 7)

% net partner counterfactual PR shares
s1_2S = S1_2S(:,1:2,:);
s1_2S(:,1,:) = S1_2S(:,1,:) + 0.5 * S1_2S(:,3,:);
s1_2S(:,2,:) = S1_2S(:,2,:) + 0.5 * S1_2S(:,3,:);
s2_2S = S2_2S(:,1:2,:);
s2_2S(:,1,:) = S2_2S(:,1,:) + 0.5 * S2_2S(:,3,:);
s2_2S(:,2,:) = S2_2S(:,2,:) + 0.5 * S2_2S(:,3,:);

table7 = table;
table7.party = {'PVEM';'PRI'};
table7.vshare_distinct_distinct = round(100 * mean(mean(S0(1:N0,3:4,:) ./ sum(S0(1:N0,:,:),2),3),1)',1);
table7.vshare_jointPRI_distinct = round(100 * mean(mean(S1(1:N0,4,:) .* s1_2S(1:N0,:,:) ./ sum(S1(1:N0,:,:),2),3),1)',1);
table7.vshare_jointPVEM_distinct = round(100 * mean(mean(S2(1:N0,3,:) .* s2_2S(1:N0,:,:) ./ sum(S2(1:N0,:,:),2),3),1)',1);
table7.vshare_distinct_jointPRI = round(100 * mean(mean(S0(N0+(1:N1),3:4,:) ./ sum(S0(N0+(1:N1),:,:),2),3),1)',1);
table7.vshare_jointPRI_jointPRI= round(100 * mean(mean(S1(N0+(1:N1),4,:) .* s1_2S(N0+(1:N1),:,:) ./ sum(S1(N0+(1:N1),:,:),2),3),1)',1);
table7.vshare_jointPVEM_jointPRI = round(100 * mean(mean(S2(N0+(1:N1),3,:) .* s2_2S(N0+(1:N1),:,:) ./ sum(S2(N0+(1:N1),:,:),2),3),1)',1);
table7.vshare_distinct_jointPVEM = round(100 * mean(mean(S0(N0+N1+(1:N2),3:4,:) ./ sum(S0(N0+N1+(1:N2),:,:),2),3),1)',1);
table7.vshare_jointPRI_jointPVEM = round(100 * mean(mean(S1(N0+N1+(1:N2),4,:) .* s1_2S(N0+N1+(1:N2),:,:) ./ sum(S1(N0+N1+(1:N2),:,:),2),3),1)',1);
table7.vshare_jointPVEM_jointPVEM = round(100 * mean(mean(S2(N0+N1+(1:N2),3,:) .* s2_2S(N0+N1+(1:N2),:,:) ./ sum(S2(N0+N1+(1:N2),:,:),2),3),1)',1);

disp(' ')
disp('for Table 7:')
disp(table7)


%% aggregate counterfactual election results (figure 3)

% optimal total-coalition configuration
[~,Yd] = max([X_PVEM * theta + ES1,X_PRI * theta + ES2],[],2);
Yd = Yd(1:N0,1);

FPTP0 = zeros(NS,5); FPTP12 = FPTP0; PR0 = FPTP0; PR12 = FPTP0; shares0 = FPTP0; shares12 = FPTP0;
for ns=1:NS

    % separately everywhere
    votes = e12.registered .* S0(:,:,ns);
    votes(1:N0,:) = [e12.MP(1:N0),e12.NA(1:N0),e12.PVEM(1:N0),e12.PRI(1:N0),e12.PAN(1:N0)]; % observed votes where PRI,PVEM ran independently
    [~,winner] = max(votes,[],2);
        % direct seats:
    FPTP0(ns,:) = [sum(winner==1),sum(winner==2),sum(winner==3),sum(winner==4),sum(winner==5)];
        % PR assignment:
    nq = sum(sum(votes)) / 200;
    PR0(ns,:) = floor(sum(votes) / nq);
    remSeats = 200 - sum(PR0(ns,:));
    if remSeats > 0
        remVotes = sum(votes) - nq * PR0(ns,:);
        [~,priority] = sort(remVotes,'descend');
        for k = 1:remSeats
            PR0(ns,priority(k)) = PR0(ns,priority(k)) + 1;
        end
    end
        % over-representation restriction:
    OR = find(FPTP0(ns,:) + PR0(ns,:) > floor(500 * (sum(votes) / sum(sum(votes)) + 0.08)));
    if ~isempty(OR)
        PR0(ns,OR) = floor(500 * (sum(votes(:,OR)) / sum(sum(votes)) + 0.08)) - FPTP0(ns,OR);
        UR = setdiff(1:5,OR);
        nq = sum(sum(votes(:,UR))) / (200 - sum(PR0(ns,OR)));
        PR0(ns,UR) = floor(sum(votes(:,UR)) / nq);
        remSeats = 200 - sum(PR0(ns,:));
        if remSeats > 0
            remVotes = sum(votes(:,UR)) - nq * PR0(ns,UR);
            [~,priority] = sort(remVotes,'descend');
            for k = 1:remSeats
                PR0(ns,UR(priority(k))) = PR0(ns,UR(priority(k))) + 1;
            end
        end
    end
    shares0(ns,:) = sum(votes) / sum(sum(votes));

    % together everywhere
    votes(find(Yd==1),:) = e12.registered(find(Yd==1)) .* [S1(find(Yd==1),1:2,ns), ...
        S1(find(Yd==1),4,ns) .* s1_2S(find(Yd==1),:,ns),S1(find(Yd==1),5,ns)]; %#ok<*FNDSB> 
    votes(find(Yd==2),:) = e12.registered(find(Yd==2)) .* [S2(find(Yd==2),1:2,ns), ...
        S2(find(Yd==2),3,ns) .* s2_2S(find(Yd==2),:,ns),S2(find(Yd==2),5,ns)];
    votes(N0+1:end,:) = [e12.MP(N0+1:end),e12.NA(N0+1:end), ...
        e12.PVEM(N0+1:end) + 0.5 * e12.CM(N0+1:end), ...
        e12.PRI(N0+1:end) + 0.5 * e12.CM(N0+1:end),e12.PAN(N0+1:end)]; % observed votes where PRI,PVEM ran together
    [~,winner([find(Yd==1);N0+(1:N1)'])] = max([votes([find(Yd==1);N0+(1:N1)'],1:2), ...
        zeros(sum(Yd==1)+N1,1),sum(votes([find(Yd==1);N0+(1:N1)'],3:4),2), ...
        votes([find(Yd==1);N0+(1:N1)'],5)],[],2);
    [~,winner([find(Yd==2);N0+N1+(1:N2)'])] = max([votes([find(Yd==2);N0+N1+(1:N2)'],1:2), ...
        sum(votes([find(Yd==2);N0+N1+(1:N2)'],3:4),2),zeros(sum(Yd==2)+N2,1), ...
        votes([find(Yd==2);N0+N1+(1:N2)'],5)],[],2);
        % direct seats:
    FPTP12(ns,:) = [sum(winner==1),sum(winner==2),sum(winner==3),sum(winner==4),sum(winner==5)];
        % PR assignment:
    nq = sum(sum(votes)) / 200;
    PR12(ns,:) = floor(sum(votes) / nq);
    remSeats = 200 - sum(PR12(ns,:));
    if remSeats > 0
        remVotes = sum(votes) - nq * PR12(ns,:);
        [~,priority] = sort(remVotes,'descend');
        for k = 1:remSeats
            PR12(ns,priority(k)) = PR12(ns,priority(k)) + 1;
        end
    end
        % over-representation restriction:
    OR = find(FPTP12(ns,:) + PR12(ns,:) > floor(500 * (sum(votes) / sum(sum(votes)) + 0.08)));
    if ~isempty(OR)
        PR12(ns,OR) = floor(500 * (sum(votes(:,OR)) / sum(sum(votes)) + 0.08)) - FPTP12(ns,OR);
        UR = setdiff(1:5,OR);
        nq = sum(sum(votes(:,UR))) / (200 - sum(PR12(ns,OR)));
        PR12(ns,UR) = floor(sum(votes(:,UR)) / nq);
        remSeats = 200 - sum(PR12(ns,:));
        if remSeats > 0
            remVotes = sum(votes(:,UR)) - nq * PR12(ns,UR);
            [~,priority] = sort(remVotes,'descend');
            for k = 1:remSeats
                PR12(ns,UR(priority(k))) = PR12(ns,UR(priority(k))) + 1;
            end
        end
    end
    shares12(ns,:) = sum(votes) / sum(sum(votes));

end

% observed vote shares
votes = [e12.MP,e12.NA,e12.PVEM + 0.5 * e12.CM,e12.PRI + 0.5 * e12.CM,e12.PAN];
shares = sum(votes) / sum(sum(votes));

figure3 = table;
figure3.party = {'MP';'NA';'PVEM';'PRI';'PAN'};
figure3.baseline_vshares = round(100 * mean(shares0)',1);
figure3.change_vshare_partial = round(100 * (shares - mean(shares0))',2);
figure3.change_vshare_total = round(100 * (mean(shares12) - mean(shares0))',2);
figure3.baseline_FPTP = round(mean(FPTP0))';
figure3.change_FPTP_partial= [71,0,19,158,52]' - round(mean(FPTP0))';
figure3.change_FPTP_total = round(mean(FPTP12))' - round(mean(FPTP0))';
figure3.baseline_PR = round(mean(PR0))';
figure3.change_PR_partial= [64,10,15,49,62]' - round(mean(PR0))';
figure3.change_PR_total = round(mean(PR12))' - round(mean(PR0))';

disp(' ')
disp('for Figure 3:')
disp(figure3)


%%

clearvars -except Yd shares ...
    FPTP0 FPTP12 PR0 PR12 shares0 shares12 ...
    table7 figure3

save('counterfactuals.mat')



